/**
  AES encryption/decryption demo program using OpenSSL EVP apis
  gcc -Wall openssl_aes.c -lcrypto

  this is public domain code. 

  Saju Pillai (saju.pillai@gmail.com)
**/


// Francesco Feltrin, 2010: transformation and integration in C++
#include "stdafx.h"

#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include "encrypt_aes.hpp"
#include "token.hpp"

/*
 * Create an 256 bit key and IV using the supplied key_data. salt can be added for taste.
*/

int encrypt_aes::aes_encrypt_pass(unsigned char *key_data, int key_data_len, unsigned char *salt, int nrounds)
{ 
  /*
   * Gen key & IV for AES 256 CBC mode. A SHA1 digest is used to hash the supplied key material.
   * nrounds is the number of times the we hash the material. More rounds are more secure but
   * slower.
   */
  int i = EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha1(), salt, key_data, key_data_len, nrounds, key, iv);
  if (i != encrypt_aes::KEY_BLOCK_SIZE) {
    fprintf(stderr, "Key size is %d bits - should be 256 bits\n", i);
    return -1;
  }

  return 0;
}

/*
 * Fills in the encryption and decryption ctx objects and returns 0 on success
 **/

int encrypt_aes::aes_encrypt_init()
{
  
  EVP_CIPHER_CTX_init(&en);
  EVP_EncryptInit_ex(&en, EVP_aes_256_cbc(), NULL, key, iv);
  
  return 0;
}

int encrypt_aes::aes_decrypt_init()
{
  
  EVP_CIPHER_CTX_init(&de);
  EVP_DecryptInit_ex(&de, EVP_aes_256_cbc(), NULL, key, iv);
  
  return 0;
}


/*
 * Encrypt *len bytes of data
 * All data going in & out is considered binary (unsigned char[])
 */

unsigned char *encrypt_aes::aes_encrypt(EVP_CIPHER_CTX *e, unsigned char *plaintext, int *len)
{
  /* max ciphertext len for a n bytes of plaintext is n + AES_BLOCK_SIZE -1 bytes */
  int c_len = *len + encrypt_aes::AES_BLOCK_SIZE, f_len = 0;
  unsigned char *ciphertext = static_cast<unsigned char*>(malloc(c_len));

  /* allows reusing of 'e' for multiple encryption cycles */
  EVP_EncryptInit_ex(e, NULL, NULL, NULL, NULL);

  /* update ciphertext, c_len is filled with the length of ciphertext generated,
    *len is the size of plaintext in bytes */
  EVP_EncryptUpdate(e, ciphertext, &c_len, plaintext, *len);

  /* update ciphertext with the final remaining bytes */
  EVP_EncryptFinal_ex(e, ciphertext+c_len, &f_len);

  *len = c_len + f_len;
  return ciphertext;
}

/*
 * Decrypt *len bytes of ciphertext
 */
unsigned char *encrypt_aes::aes_decrypt(EVP_CIPHER_CTX *e, unsigned char *ciphertext, int *len)
{
  /* plaintext will always be equal to or lesser than length of ciphertext*/
  int p_len = *len, f_len = 0;
  unsigned char *plaintext = static_cast<unsigned char*>(malloc(p_len));
  
  EVP_DecryptInit_ex(e, NULL, NULL, NULL, NULL);
  EVP_DecryptUpdate(e, plaintext, &p_len, ciphertext, *len);
  EVP_DecryptFinal_ex(e, plaintext+p_len, &f_len);

  *len = p_len + f_len;
  return plaintext;
}


//int encrypt_aes::encrypt(const char* _input, int *_inputlen)
int encrypt_aes::encrypt()
{
	 /* the key_data is read from the argument list */
   int inputlen = strlen(input)+1;
		
	/* 8 bytes to salt the key_data during key generation. This is an example of
     	compiled in salt. We just read the bit pattern created by these two 4 byte 
     	integers on the stack as 64 bits of contigous salt material - 
     	of course this only works if sizeof(int) >= 4 */
  	
	const unsigned int salt[] = {12345, 54321};
	 
   /* gen key and iv. init the cipher ctx object */
  if (aes_encrypt_pass(key_data, key_data_len, (unsigned char *)&salt) || aes_encrypt_init()) {
    fprintf(stderr, "Couldn't initialize AES cipher\n");
    return -1;
  }

  /* encrypt and decrypt each input string and compare with the original */
       
    
    /* The enc/dec functions deal with binary data and not C strings. strlen() will 
       return length of the string without counting the '\0' string marker. We always
       pass in the marker byte to the encrypt/decrypt functions so that after decryption 
       we end up with a legal C string */
    
    ciphertext = aes_encrypt(&en, (unsigned char*)input, &inputlen);
    ciphertextlen = inputlen;

  EVP_CIPHER_CTX_cleanup(&en);
 
  return 0;
}
  
const std::string encrypt_aes::decrypt()
{	   
  /* init the cipher ctx object */
  if (aes_decrypt_init()) {
    fprintf(stderr, "Couldn't initialize AES cipher\n");
    return NULL;
  }

  /* encrypt and decrypt each input string and compare with the original */
    plaintext = (char *)aes_decrypt(&de, ciphertext, &ciphertextlen);
	/*
#ifdef _DEBUG
    fprintf(stderr, "%s\n", plaintext);        
#endif
	*/
 
  EVP_CIPHER_CTX_cleanup(&de);
  return plaintext;
}


encrypt_aes::~encrypt_aes()
{
	if(plaintext) free(plaintext); plaintext = NULL;
	if(ciphertext) free(ciphertext); ciphertext = NULL;
}	

encrypt_aes::encrypt_aes(const char* _input, const char* _passwd): 
	ciphertext(NULL), plaintext(NULL), input(_input) 
{
	key_data = (unsigned char*) _passwd;
	key_data_len = strlen(_passwd);
}

void encrypt_aes::pack(const std::string& _filename, unsigned char* _str, int outlen)
{
	FILE* out = NULL;

	if(_filename.length())
		out = fopen(_filename.c_str(), "wb");

	if((out==NULL))
	{
		fprintf(stderr, "Cannot write into file: '%s'\n", _filename.c_str());
	} 
	else	
	{
		fwrite(_str, 1, outlen, out); 
		fclose(out);
	}	
}

void encrypt_aes::save(const std::string& _cipher_filename, 
					   const std::string& _key_filename, 
					   const std::string& _iv_filename)
{
	pack(_cipher_filename, ciphertext, ciphertextlen);
	pack(_key_filename, key, KEY_BLOCK_SIZE);
	pack(_iv_filename, iv, KEY_BLOCK_SIZE);
}

int encrypt_aes::unpack(const std::string& _filename, unsigned char** _value)
{
	FILE *_file = fopen(_filename.c_str(),"rb");
	unsigned char* _ptr = *_value = (unsigned char*) malloc(sizeof(unsigned char)*KEY_BLOCK_SIZE);
	unsigned short i;

	if(_file == NULL)
	 {
  		fprintf(stderr, "\nError opening file: '%s'\n", _filename.c_str());
  		return -1;
 	}

//#ifdef _DEBUG
 	fprintf(stderr, "\n Public key file opened\n");
//#endif

		//load the key
	for(i=0;!feof(_file) && _ptr;) 
	{
		do
		{
 			fread(&_ptr[i++],sizeof(unsigned char),1,_file);			
		} while ((i % KEY_BLOCK_SIZE) && !feof(_file));

		_ptr =  (unsigned char*) realloc(_ptr, sizeof(unsigned char)*KEY_BLOCK_SIZE*(i/KEY_BLOCK_SIZE+1));
	}
	
	_ptr = (unsigned char*) realloc(_ptr, sizeof(unsigned char)*(--i));

	*_value = _ptr;

 	fclose(_file);
	return i;
}

encrypt_aes::encrypt_aes(): ciphertext(NULL), ciphertextlen(0), plaintext(NULL) {}

bool encrypt_aes::read(const std::string& _ciphertext_file, 
					   const std::string& _key_file, const std::string& _iv_file)
{
	unsigned char *_ptr;
	int j(0), l;

	l = ciphertextlen = unpack(_ciphertext_file, &ciphertext);
	
	if((ciphertextlen >0) &&
		(j=unpack(_key_file, &_ptr)) == KEY_BLOCK_SIZE)
	{
		memcpy(key, _ptr, sizeof(unsigned char)*KEY_BLOCK_SIZE);
		free(_ptr); _ptr =NULL;
	}
	l*=j;

	if((l > 0) && 
		(j=unpack(_iv_file, &_ptr)) == KEY_BLOCK_SIZE)
	{
		memcpy(iv, _ptr, sizeof(unsigned char)*KEY_BLOCK_SIZE);
		free(_ptr); _ptr =NULL;
	}

	return (l*j > 0);
}

bool encrypt_aes::check_decrypt(const char* _input)
{
   bool b(false);

   if(_input)
   {
   	if ((b= strncmp(plaintext, _input, strlen(_input)+1)))
      		fprintf(stderr, "FAIL: enc/dec failed for \"%s\"\n", _input);
    	else
      		fprintf(stderr, "OK: enc/dec ok for \"%s\"\n", plaintext);
   }

    return !b;
}


void encrypt_aes::encrypt_into(const char* _input, const char* _passwd, const std::string& _cipher_file, const std::string& _key_file, const std::string& _iv_file)
{
#ifdef _DEBUG
	std::cout << "password to cipher = " << _input << std::endl;
	std::cout << "ciphering password = " << _passwd << std::endl;
#endif
	encrypt_aes aes((const char*)_input, _passwd);
	aes.encrypt();
	aes.save(_cipher_file, _key_file, _iv_file);
}


const std::string encrypt_aes::decrypt_from(const std::string& _cipher_file, const std::string& _key_file, 
		const std::string& _iv_file, const std::vector<unsigned short>& _substring_indices)
{
	encrypt_aes aes;
	std::string decrypted_str;

	if(aes.read(_cipher_file, _key_file, _iv_file) > 0)
		decrypted_str = aes.decrypt();		

	if(_substring_indices.capacity() && decrypted_str.length())
	{
    std::string str;

    for(unsigned short k =0; k < _substring_indices.capacity(); k++)
    {
       if(_substring_indices[k] < decrypted_str.length() && _substring_indices[k] >=0 )
  	     str.push_back(decrypted_str[_substring_indices[k]]);
    }     
    return str;
	}

	return decrypted_str;
}

